import torch
import torch.nn as nn


class EmotionLSTM(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers=2):
        super(EmotionLSTM, self).__init__()

        # LSTM层
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True, dropout=0.2)

        # 全连接层
        self.fc = nn.Linear(hidden_dim, output_dim)

        # Softmax层，输出每个类的概率
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        # x shape: (batch_size, sequence_length, input_dim)

        lstm_out, (h_n, c_n) = self.lstm(x)  # lstm_out shape: (batch_size, sequence_length, hidden_dim)

        # 只取LSTM最后时刻的输出
        last_hidden_state = lstm_out[:, -1, :]  # shape: (batch_size, hidden_dim)

        # 全连接层
        out = self.fc(last_hidden_state)  # shape: (batch_size, output_dim)

        return self.softmax(out)  # 输出情绪的概率
